import jax
import jax.numpy as jnp
import equinox as eqx
from jaxtyping import Array, Float, Int, Bool, PyTree
from .modules.utils import forward_pass, get_dts
from .modules.transformer import MultiheadAttention
from .modules.temporal_enc import TemporalEnc

class CTTransformer(eqx.Module):
    temporal_enc: TemporalEnc
    stack: list
    out_linear: eqx.Module

    def __init__(self, hdim: int, embed_dim: int, nlayers: int, nhead: int, dropout_rate: float, key: Array):
        keys = jax.random.split(key, nhead*nlayers+1)
        self.temporal_enc = TemporalEnc(embed_dim)
        self.stack = [
            [MultiheadAttention(2*embed_dim, 1, hdim, hdim, dropout_rate, keys[i*nlayers+j]) for j in range(nlayers)]
            for i in range(nhead)
        ]
        self.out_linear = eqx.nn.Linear(2*embed_dim, embed_dim, key=keys[-1])

    @eqx.filter_jit
    def __call__(self, query_t: Float[Array, "T_"], ts: Float[Array, "T"], embeds: Float[Array, "T embed_dim"], 
                                                                    mask: Bool[Array, "T_ T"], key: Array):
        nhead = len(self.stack)
        nlayer = len(self.stack[0])
        t_embed = self.temporal_enc(ts) # (T, embed_dim)
        query_t_embed = self.temporal_enc(query_t) # (T_, embed_dim)
        T = ts.shape[0]
        T_ = query_t.shape[0]
        sub_mask = jnp.tril(jnp.ones((T, T), dtype=bool)) # (T, T)
        upper_mask = jnp.concatenate((sub_mask, jnp.zeros((T, T_), dtype=bool)), 1) # (T, T+T_)
        lower_mask = jnp.concatenate((mask, jnp.eye(T_, dtype=bool)), 1) # (T_, T+T_)
        full_mask = jnp.concatenate((upper_mask, lower_mask), 0) # (T+T_, T+T_)
        hs = list()
        X = jnp.concatenate((embeds, t_embed), -1) # (T, 2*embed_dim)
        for head_i in range(nhead):
            h = jnp.zeros((T_, embeds.shape[-1]))
            for layer_j in range(nlayer):
                X_ = jnp.concatenate((h, query_t_embed), -1) # (T_, 2*embed_dim)
                combined_input = jnp.concatenate((X, X_), 0) # (T+T_, 2*embed_dim)
                cur_key, key = jax.random.split(key)
                output = self.stack[head_i][layer_j](combined_input, combined_input, combined_input, full_mask, cur_key)
                output = jax.vmap(self.out_linear)(output) # (T+T_, embed_dim)
                h = h + output[T:]
                X = jnp.concatenate((output[:T], t_embed), -1)
            hs.append(h)
        h = jnp.concatenate(hs, -1) # (T_, nhead*embed_dim)
        return h
   
class AttNHP(eqx.Module):
    enc: CTTransformer
    intensity_fn: list
    embed: eqx.Module

    def __init__(self, hdim: int, num_types: int, embed_dim: int, nlayers: int, nhead: int, key: Array):
        key1, key2, key3 = jax.random.split(key, 3)
        self.enc = CTTransformer(hdim, embed_dim, nlayers, nhead, .1, key1)
        self.embed = eqx.nn.Embedding(num_types, embed_dim, key=key2)
        self.intensity_fn = [
            eqx.nn.Linear(nhead*embed_dim, num_types, key=key3),
            jax.nn.softplus
        ]

    @eqx.filter_jit
    def _get_intensities(self, query_t: Float[Array, "T_"], ts: Float[Array, "T"], embeds: Int[Array, "T embed_dim"], mask: Bool[Array, "T_ T"], key: Array):
        h = self.enc(query_t, ts, embeds, mask, key) # (T_, nhead*embed_dim)
        intensities = jax.vmap(forward_pass, (None, 0))(self.intensity_fn, h) # (T_, num_types, )
        return intensities

    @eqx.filter_jit
    def _lls(self, ts_arg: Float[Array, "T"], marks_arg: Int[Array, "T"], mask_arg: Bool[Array, "T"], key_arg):
        T = ts_arg.shape[0]
        dts = get_dts(ts_arg)
        dts = jnp.where(mask_arg, dts, 0.)
        embeds = jax.vmap(self.embed)(marks_arg)
        sample_ratios = jnp.linspace(0., 1., 100) # (100, )
        sampled_dts = sample_ratios[:, None] * dts[None, 1:]  # (100, T-1)
        sampled_ts = ts_arg[:-1] + sampled_dts # (100, T-1)
        attn_mask = jnp.tril(jnp.ones((T-1, T-1), dtype=bool)) # (T-1, T-1)

        intensities = jax.vmap(self._get_intensities, (0, None, None, None, 0))(sampled_ts, ts_arg[:-1], embeds[:-1], attn_mask, jax.random.split(key_arg, 100)) # (T-1, num_types)

        event_intensities = intensities[-1] # (T-1, num_types)
        non_intensities = intensities[:-1].sum(-1) # (99, T-1)
        Lambda = non_intensities.mean(0) * dts[1:] # (T-1, )
        event_intensity = jnp.take_along_axis(event_intensities, marks_arg[1:, None], -1).squeeze(-1) # (T-1,)
        ll = jnp.log(event_intensity+1e-6) - Lambda # (T-1, )
        mark_dist = event_intensity / (event_intensities.sum(-1)+1e-6) # (T-1)
        mark_ll = jnp.log(mark_dist+1e-6)
        ll = jnp.where(mask_arg[1:], ll, 0.)
        mark_ll = jnp.where(mask_arg[1:], mark_ll, 0.)
        dt_ll = ll - mark_ll
        return ll, dt_ll, mark_ll # (T-1, )

    def __call__(self, ts: Float[Array, "T"], marks: Int[Array, "T"], mask: Bool[Array, "T"], key: Array):
        ll, dt_ll, mark_ll = self._lls(ts, marks, mask, key)
        ll = ll.sum()
        dt_ll = dt_ll.sum()
        mark_ll = mark_ll.sum()
        return ll, dt_ll, mark_ll
    
    def intensity_at(self, h: Float[Array, "hdim"], dt: Float[Array, ""]):
        raise NotImplementedError()